package com.framework.buss.ai.controller;

import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.util.ObjectUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;
import java.util.UUID;

/**
 * @author zzg
 * @date 2025-2-24 15:26
 */
@Slf4j
@RestController
@RequestMapping("/ai")
@Api(value = "ai", tags = "ai问答")
public class AiController {
    private static final String END = "stop";
    @Autowired
    ClientV4 client;

    @PostMapping(value = "/completion", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @ApiOperation(value = "对话")
    public SseEmitter completion(@RequestBody List<ChatMessage> messages) {
        String requestId = UUID.randomUUID().toString().replace("-", "");
        SseEmitter emitter = new SseEmitter(-1L);
        // "codegeex-4"
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4Flash)
                .stream(Boolean.TRUE)

                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();

        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        if (sseModelApiResp.isSuccess()) {
            sseModelApiResp.getFlowable().subscribe(it -> {
                String end = it.getChoices().get(0).getFinishReason();
                if (!ObjectUtils.isEmpty(end) && END.equals(end)) {
                    emitter.send(SseEmitter.event().data("[DONE]").build());
                    emitter.complete();
                } else {
                    String content = it.getChoices().get(0).getDelta().getContent();
                    emitter.send(SseEmitter.event().data(content).build());
                }
            }, err -> emitter.completeWithError(err));
        }

        return emitter;
    }
}
